import sys
import matplotlib.pyplot as plt
from src.data.dataset import FontDataset
from src.configs import base_config

def main():
    dataset = FontDataset()
    print(f"样本数量: {len(dataset)}")
    
    if len(dataset) > 0:
        # 测试collate_fn
        from src.data.dataset import collate_fn
        dummy_batch = [dataset[i] for i in range(2)]
        batch = collate_fn(dummy_batch)
        
        print("\nBatch数据验证:")
        print(f"图像形状: {batch['image'].shape}")
        print(f"Token矩阵形状: {batch['tokens'].shape}")
        print(f"实际长度: {batch['seq_len'].tolist()}")
        
        # 可视化第一个样本
        plt.figure(figsize=(10,5))
        # 在可视化部分添加通道检查
        plt.subplot(1,2,1)
        if batch['image'].shape[1] == 1:  # [B,1,H,W]
            plt.imshow(batch['image'][0].permute(1,2,0).mean(dim=2).numpy(), cmap='gray')
        else:
            plt.imshow(batch['image'][0].permute(1,2,0).numpy() * 0.5 + 0.5)
        plt.title(f"Channels: {batch['image'].shape[1]}")
        plt.title("Sample Image")
        
        plt.subplot(1,2,2)
        plt.plot(batch['tokens'][0][:batch['seq_len'][0]].numpy())
        plt.title("Padded Token Sequence")
        
        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    main()